import pickle

import numpy as np
from matplotlib import pyplot as plt


def launch():
    # Load parameters data
    with open("policy_gradient/parameters_data_pg.pkl", 'rb') as f:
        params_data_pg = pickle.load(f)

    max_rew_lr_curv_pg = params_data_pg['max_rew_lr_curve']
    mean_rew_lr_curve_pg = params_data_pg['mean_rew_lr_curve']
    slow_down_lr_curve_pg = params_data_pg['slow_down_lr_curve']
    ref_discount_rews_pg = params_data_pg['ref_discount_rews']
    ref_slow_down_pg = params_data_pg['ref_slow_down']
    elapsed_time_pg = params_data_pg['elapsed_time']

    with open("A2C/parameters_data_a2c.pkl", 'rb') as f:
        params_data_a2c = pickle.load(f)

    max_rew_lr_curv_a2c = params_data_a2c['max_rew_lr_curve']
    mean_rew_lr_curve_a2c = params_data_a2c['mean_rew_lr_curve']
    slow_down_lr_curve_a2c = params_data_a2c['slow_down_lr_curve']
    ref_discount_rews_a2c = params_data_a2c['ref_discount_rews']
    ref_slow_down_a2c = params_data_a2c['ref_slow_down']
    elapsed_time_a2c = params_data_a2c['elapsed_time']


    plot_data(ref_discount_rews_pg, mean_rew_lr_curve_pg, max_rew_lr_curv_pg, elapsed_time_pg, ref_slow_down_pg,
              ref_discount_rews_a2c, mean_rew_lr_curve_a2c, max_rew_lr_curv_a2c, elapsed_time_a2c,
              ref_slow_down_a2c)


def plot_data(ref_discount_rews_pg, mean_rew_lr_curve_pg, max_rew_lr_curve_pg, elapsed_time_pg, ref_slow_down_pg,
              ref_discount_rews_a2c, mean_rew_lr_curve_a2c, max_rew_lr_curve_a2c, elapsed_time_a2c,
              ref_slow_down_a2c):
    num_colors = len(ref_discount_rews_pg) + 2
    cm = plt.get_cmap('gist_rainbow')

    fig = plt.figure(figsize=(12, 5))

    ax = fig.add_subplot(121)
    # Replace this line
    # ax.set_color_cycle([cm(1. * i / num_colors) for i in range(num_colors)])

    # With this line
    ax.set_prop_cycle('color', [cm(1. * i / num_colors) for i in range(num_colors)])

    #ax.plot(mean_rew_lr_curve_pg, linewidth=2, label='PG mean')
    # for k in ref_discount_rews_pg:
    # ax.plot(np.tile(np.average(ref_discount_rews_pg[k]), len(mean_rew_lr_curve_pg)), linewidth=2, label=k)
    ax.plot(max_rew_lr_curve_pg, linewidth=2, label='PG max')

    #ax.plot(mean_rew_lr_curve_a2c, linewidth=2, label='A2C mean')
    # for k in ref_discount_rews_a2c:
    #     ax.plot(np.tile(np.average(ref_discount_rews_a2c[k]), len(mean_rew_lr_curve_a2c)), linewidth=2, label=k)
    ax.plot(max_rew_lr_curve_a2c, linewidth=2, label='A2C max')

    plt.legend(loc=4)
    plt.xlabel("Iteration", fontsize=20)
    plt.ylabel("Discounted Total Reward", fontsize=20)

    ax = fig.add_subplot(122)
    # Replace this line
    # ax.set_color_cycle([cm(1. * i / num_colors) for i in range(num_colors)])

    # With this line
    ax.set_prop_cycle('color', [cm(1. * i / num_colors) for i in range(num_colors)])


    ax.plot(elapsed_time_pg, linewidth=2, label='PG time')
    # for k in ref_discount_rews_pg:
    #     ax.plot(np.tile(np.average(np.concatenate(ref_slow_down_pg[k])), len(slow_down_lr_curve_pg)), linewidth=2,
    #             label=k)

    ax.plot(elapsed_time_a2c, linewidth=2, label='A2C time')
    # for k in ref_discount_rews_a2c:
    #     ax.plot(np.tile(np.average(np.concatenate(ref_slow_down_a2c[k])), len(slow_down_lr_curve_a2c)), linewidth=2,
    #             label=k)

    plt.legend(loc=1)
    plt.xlabel("Iteration", fontsize=20)
    plt.ylabel("ElapsedTime", fontsize=20)
    plt.savefig('data_comparison.pdf')
    print("对比图已完成")

    # 关闭图表窗口
    plt.close()


def main():
    launch()


if __name__ == '__main__':
    main()
